import java.security.*;
import java.security.interfaces.*;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.BadPaddingException;

import java.nio.charset.Charset;

import java.math.BigInteger;
import java.io.IOException;
import java.io.FileWriter;
import java.io.File;
import java.util.Scanner; 

import org.apache.commons.codec.binary.Hex;

public class RSAEncrypt {
    public static void main(String args[]) {
        String plaintext, ciphertext;
        plaintext = "";
        ciphertext = "";
        boolean decryptOnly, readFromFile;
        readFromFile = false;
        decryptOnly = false;
        if (args.length == 1) plaintext = args[0];
        else if (args.length == 2) {
            readFromFile = true;
            if (args[0].equals("-f")) { /* read keys from files */
                plaintext = args[1];

            } else if (args[0].equals("-fd")) { /* keys from files AND read input
                                                 * as ciphertext, and decrypt only */
                ciphertext = args[1];
            }
        }
    
        if (plaintext.length() > 0 || ciphertext.length() > 0) {
            try { 
                RSACrypto rsaCrypto;
                if (!readFromFile) rsaCrypto = new RSACrypto();
                else rsaCrypto = new RSACrypto(true);

                System.out.println("RSA Encryption/Decryption");
                System.out.println("\tpublic key = " + rsaCrypto.getPublicKey());
                System.out.println("\tprivate key = " + rsaCrypto.getPrivateKey());
                System.out.println("\tmodulus = " + rsaCrypto.modulus());
                if (ciphertext.length() == 0) 
                    System.out.println("\tplaintext = " + plaintext);
                else 
                    System.out.println("\tciphertext = " + ciphertext);;

                if (ciphertext.length() == 0) {
                    ciphertext = rsaCrypto.encrypt(plaintext);
                    System.out.println("\tciphertext = " + ciphertext);
                }
                String decryptedCt = rsaCrypto.decrypt(ciphertext);
                System.out.println("\tdecrypted ciphertext = " + decryptedCt);
            } catch (Exception e) {
                e.printStackTrace();
            }
            
            System.out.println();

        } 
    }
}

class RSACrypto {
    private static final String ALGORITHM = "RSA/ECB/NoPadding";    
    private static final String ALGORITHM_KEY = "RSA";
    private static final int KEY_SIZE = 1024;
    // private static final int KEY_SIZE = 512;
    private static final String ENCODING = "UTF-8";

    private static final String PUBLIC_KEY_FILE = "publickey.txt";
    private static final String PRIVATE_KEY_FILE = "privatekey.txt";
    private static final String MODULUS_FILE = "modulus.txt";

    private KeyPair keyPair;
    private PublicKey publicKey;
    private PrivateKey privateKey;
    private BigInteger modulus;
    private Cipher cipher;    

    public RSACrypto() throws Exception {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance(ALGORITHM_KEY);
        kpg.initialize(KEY_SIZE);
        this.keyPair = kpg.generateKeyPair();
        this.publicKey = keyPair.getPublic();
        this.privateKey = keyPair.getPrivate();
        this.modulus = ((RSAPublicKey)publicKey).getModulus();
        this.cipher = Cipher.getInstance(ALGORITHM);
        writeKeyFiles();
    }

    private void writeKeyFiles() {
        try { 
            FileWriter fileWriter = new FileWriter(PUBLIC_KEY_FILE);
            fileWriter.write(getPublicKey().toString());
            fileWriter.close();
 
            fileWriter = new FileWriter(PRIVATE_KEY_FILE);
            fileWriter.write(getPrivateKey().toString());
            fileWriter.close();

            fileWriter = new FileWriter(MODULUS_FILE);
            fileWriter.write(modulus.toString());
            fileWriter.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }    

    private BigInteger[] readKeyFiles() {
        BigInteger[] keyComponents = new BigInteger[3];
        try { 
            Scanner fileReader = new Scanner(new File(PUBLIC_KEY_FILE));
            keyComponents[0] = new BigInteger(fileReader.nextLine()); 
            fileReader.close();

            fileReader = new Scanner(new File(PRIVATE_KEY_FILE));
            keyComponents[1] = new BigInteger(fileReader.nextLine());
            fileReader.close();

            fileReader = new Scanner(new File(MODULUS_FILE));
            keyComponents[2] = new BigInteger(fileReader.nextLine());
            fileReader.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return keyComponents;
    }

    public RSACrypto(boolean keysFromFiles) throws Exception {
        BigInteger[] keyComponents = readKeyFiles();
        BigInteger publicExp, privateExp, modulus;
        publicExp = keyComponents[0];
        privateExp = keyComponents[1];
        modulus = keyComponents[2];

        RSAPublicKeySpec rsaKeySpec = new RSAPublicKeySpec(modulus, publicExp);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_KEY);
        this.publicKey = keyFactory.generatePublic(rsaKeySpec);

        RSAPrivateKeySpec rsaPrivKeySpec = new RSAPrivateKeySpec(modulus, privateExp);
        this.privateKey = keyFactory.generatePrivate(rsaPrivKeySpec);

        // this.privateKey = null;
        this.modulus = ((RSAPublicKey)publicKey).getModulus();
        this.cipher = Cipher.getInstance(ALGORITHM);
    }


    public String encrypt(String plaintext) throws Exception {
        // cipher.init(Cipher.ENCRYPT_MODE, keyPair.getPublic());
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] ptBytes = plaintext.getBytes(ENCODING);
        // byte[] ptBytes = Base64.decode(ptBytes);
        byte[] ctBytes = blockCipher(ptBytes, Cipher.ENCRYPT_MODE);
         
        char[] ctBytesAsChars = Hex.encodeHex(ctBytes);
        // char[] ctBytesAsChars = 
        return new String(ctBytesAsChars);
        //return Base64.encodeBytes(ctBytes);
    }

    public String decrypt(String ciphertext) throws Exception {
        // cipher.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] ctBytes = Hex.decodeHex(ciphertext.toCharArray());
        // byte[] ctBytes = Base64.decode(ciphertext);
        byte[] decryptedBytes = blockCipher(ctBytes, Cipher.DECRYPT_MODE);
        return new String(decryptedBytes, ENCODING);
        // return new String(Base64.encodeBytesToBytes(decryptedBytes), ENCODING);
    }
    
    private byte[] blockCipher(byte[] bytes, int mode) throws IllegalBlockSizeException, BadPaddingException{
        byte[] scrambled = new byte[0];

        byte[] toReturn = new byte[0];
        int length = (mode == Cipher.ENCRYPT_MODE)? 100 : 128;
        byte[] buffer = new byte[length];

        for (int i=0; i< bytes.length; i++){

            if ((i > 0) && (i % length == 0)){
                scrambled = cipher.doFinal(buffer);
                toReturn = append(toReturn,scrambled);
                int newlength = length;

                if (i + length > bytes.length) {
                     newlength = bytes.length - i;
                }
                buffer = new byte[newlength];
            }
            buffer[i%length] = bytes[i];
        }

        scrambled = cipher.doFinal(buffer);

        toReturn = append(toReturn,scrambled);

        return toReturn;
    }
    
    private byte[] append(byte[] prefix, byte[] suffix){
		byte[] toReturn = new byte[prefix.length + suffix.length];
		for (int i=0; i< prefix.length; i++) toReturn[i] = prefix[i];
		for (int i=0; i< suffix.length; i++) toReturn[i+prefix.length] = suffix[i];
		return toReturn;
	}

    public BigInteger getPublicKey() {
        return ((RSAPublicKey)publicKey).getPublicExponent();
    }    
        
    public BigInteger getPrivateKey() {
        return ((RSAPrivateKey)privateKey).getPrivateExponent();
    }

    public BigInteger modulus() {
        return modulus;
    }
}

